import pybullet as p
import pybullet_data
import math
import numpy as np
import time
from typing import List, Tuple, Dict, Optional


class InteractiveRobotTrajectorySimulator:
    """
    机械臂桁轨仿真与碰撞检测接口
    """

    def __init__(self, gui_mode: bool = True, collision_threshold: float = 0.01):
        """
        初始化仿真环境

        Args:
            gui_mode: 是否显示GUI界面
            collision_threshold: 碰撞检测阈值
        """
        # 连接物理引擎
        if gui_mode:
            self.physics_client = p.connect(p.GUI)
        else:
            self.physics_client = p.connect(p.DIRECT)

        p.setAdditionalSearchPath(pybullet_data.getDataPath())

        # 碰撞检测参数（固定值）
        self.collision_detection_enabled = True
        self.collision_threshold = collision_threshold

        # 初始化组件
        self.obstacle = None  # 桁轨（保持原始命名）
        self.robot_id = None
        self.box_body = None  #单独添加的滑动横梁box
        self.end_effector = None   #末端执行器
        self.mount_constraint = None
        self.gantry_prismatic_joints = []
        self.robot_joints = []
        self.arm_cur_angle = []
        self.collision_info = []  #碰撞信息列表
        '''
        ==================collision_info[]列表参数说明==================
            contactFlag,         # 内部标志位（通常可忽略）
            bodyUniqueIdA,       # 刚体A的ID
            bodyUniqueIdB,       # 刚体B的ID
            linkIndexA,         # 刚体A的链接索引
            linkIndexB,         # 刚体B的链接索引
            positionOnA,        # 在刚体A上的接触位置[x,y,z]
            positionOnB,        # 在刚体B上的接触位置[x,y,z]
            contactNormalOnB,   # 刚体B上的接触法向量[x,y,z]
            contactDistance,    # 接触距离（正值为穿透深度）
            normalForce,        # 法向力大小（需要开启接触力计算）
            lateralFriction1,   # 第一横向摩擦力
            lateralFriction2,   # 第二横向摩擦力
            lateralFrictionDir1,# 第一横向摩擦方向
            lateralFrictionDir2 # 第二横向摩擦方向
        '''

        self.x_cur_pos = 0
        self.y_cur_pos = 0


        self._setup_scene_original()

    def _setup_scene_original(self):
        """设置仿真场景"""
        # 加载桁轨
        self.obstacle = p.loadURDF("./333333/urdf/3333332.urdf", [1.2885, 2.986, 0], [0, 0, 0, 1], useFixedBase=True)

        # 机械臂旋转180度（倒装）- 绕Y轴旋转180度
        robot_orientation = p.getQuaternionFromEuler([0, math.pi, 0])  # 绕Y轴旋转180度

        # 加载机械臂（倒装）
        self.robot_id = p.loadURDF("elfin10_l.urdf", [0, 0, 2], robot_orientation, flags=p.URDF_USE_SELF_COLLISION)
        #加载末端执行器

        self.end_effector = p.loadURDF("./Endeffector/urdf/Endeffector.urdf", [0, 0, 0], [0, 0, 0, 1])

#======================================建立机械臂末端与x滑动关节之间的约束===========================================
        # 获取桁轨滑动关节
        self.gantry_prismatic_joints = []
        num_gantry_joints = p.getNumJoints(self.obstacle)
        for i in range(num_gantry_joints):
            joint_info = p.getJointInfo(self.obstacle, i)
            if joint_info[2] == p.JOINT_PRISMATIC:
                self.gantry_prismatic_joints.append(i)
        print("桁轨滑动关节索引：", self.gantry_prismatic_joints)

        # 创建约束
        self.mount_constraint = p.createConstraint(
            self.obstacle, self.gantry_prismatic_joints[1], self.robot_id, -1,
            p.JOINT_FIXED, [0, 0, 0], [0, 0, 0],
            [0, 0, 0], [0, 0, 0, 1], [0, 0, 0, 1]
        )
#=====================================添加box到y横梁，充当y横梁的碰撞体================================================

        box_half_extents = [0.231, 1.6, 0.09]  # 半轴长度
        # 创建碰撞形状
        collision_shape = p.createCollisionShape(p.GEOM_BOX, halfExtents=box_half_extents)
        # 创建视觉形状
        visual_shape = p.createVisualShape(
            shapeType=p.GEOM_BOX,
            halfExtents=box_half_extents,
            rgbaColor=[0.75294, 0.75294, 0.75294, 1]
        )
        # 选择要附加的关节（这里使用第一个滑动关节）
        target_joint_index = self.gantry_prismatic_joints[0]

        # 计算BOX相对于关节的位置偏移
        # 需要根据实际关节位置调整这些值
        box_offset_pos = [0, 0, 0]  # 相对于关节的局部位置
        box_offset_orn = p.getQuaternionFromEuler([0, 0, 0])  # 无旋转

        # 创建BOX并约束到关节
        self.box_body = p.createMultiBody(baseMass=1,  # 0表示静态物体
                                          baseCollisionShapeIndex=collision_shape,
                                          baseVisualShapeIndex=visual_shape,
                                          basePosition=[0, 0, 3.5],  # 初始位置会被约束覆盖
                                          baseOrientation=[0.7071, 0.7071, 0, 0])

        # 创建约束将BOX附加到关节
        box_constraint = p.createConstraint(self.obstacle, target_joint_index,
                                            self.box_body, -1,
                                            p.JOINT_FIXED,
                                            jointAxis=[0, 0, 0],
                                            parentFramePosition=box_offset_pos,
                                            childFramePosition=[-0.151, 0, 0],
                                            parentFrameOrientation=box_offset_orn,
                                            childFrameOrientation=[0, 0, 0, 1])

#============================================添加末端执行器与机械臂末端之间的约束==============================================

        # 获取机械臂关节

        end_effector_orientation = p.getQuaternionFromEuler([math.pi / 2, math.pi / 8, 0])
        end_effector_constraint = p.createConstraint(self.robot_id, 6, self.end_effector, -1,
                                                     p.JOINT_FIXED,
                                                     jointAxis=[0, 0, 0],
                                                     parentFramePosition=[0, 0, 0],
                                                     childFramePosition=[-0.065, 0.25, 0],
                                                     parentFrameOrientation=[0, 0, 0, 1],
                                                     childFrameOrientation=end_effector_orientation)

#====================================================================================================================
        self._get_robot_joints()
        print(f"机械臂关节数量: {len(self.robot_joints)}")

    def _get_robot_joints(self):
        """获取机械臂的旋转关节"""
        self.robot_joints = []
        num_joints = p.getNumJoints(self.robot_id)
        for i in range(num_joints):
            joint_info = p.getJointInfo(self.robot_id, i)
            if joint_info[2] == p.JOINT_REVOLUTE:
                self.robot_joints.append(i)


    def get_end_effector_pose(self) -> Tuple[List[float], List[float]]:
        """获取末端执行器的位置和姿态"""
        # 获取末端链接的状态
        link_state = p.getBasePositionAndOrientation(self.end_effector)
        position = link_state[0]
        orientation = link_state[1]
        return position, orientation

    def get_prismatic_joints_pose(self) -> Tuple[float, float]:
        """获取滑动关节的位置"""
        # 获取滑动关节的位置
        joint_y_state = p.getJointState(self.obstacle, self.gantry_prismatic_joints[0])
        joint_x_state = p.getJointState(self.obstacle, self.gantry_prismatic_joints[1])
        y_position = joint_y_state[0]
        x_position = joint_x_state[0]
        return x_position, y_position

    def get_arm_joints_angle(self) -> List[float] :
        cur_angle = []
        for i in self.robot_joints:
            joint_state = p.getJointState(self.robot_id, i)
            cur_angle.append(round(joint_state[0],4))
        return cur_angle

    def check_collision(self) -> bool:
        """检查是否发生碰撞"""
        # 检查机械臂与桁轨的碰撞
        contact_points01 = p.getContactPoints(self.robot_id)  #检测机械臂和任何其他物体的碰撞
        contact_points02 = p.getContactPoints(self.end_effector)

        # 过滤掉约束连接点的接触
        collision_detected = False
        for contact in contact_points01:
            # 如果接触距离小于阈值，认为是碰撞
            if contact[8] < self.collision_threshold:  # contact[8] 是接触距离
                collision_detected = True
                self.collision_info = contact_points01
                break
        for contact in contact_points02:
            # 如果接触距离小于阈值，认为是碰撞
            if contact[8] < self.collision_threshold:  # contact[8] 是接触距离
                collision_detected = True
                self.collision_info = contact_points02
                break
        return collision_detected


    def simulate_trajectory(self, trajectory: List[Tuple[List[float], float, float]],
                            gui_mode: bool = True,
                            step_time: float = 0,
                            real_time: bool = True) -> Dict:
        """
        播放预定义轨迹

        Args:
            trajectory: 轨迹列表，每个元素为(joint_angles, x_pos, y_pos)
            gui_mode : 是否可视化
            step_time: 每步的时间间隔
            real_time: 是否实时显示

        Returns:
            仿真结果字典
        """
        results = {
            'collision_position': [],    # 碰撞时机械臂末端位置
            'collision_detected': False, # 是否发生碰撞
            'collision_step': -1,        # 碰撞发生的步数(发生碰撞时的轨迹点索引)
            'collision_info': []
        }

        print(f"开始播放轨迹，共 {len(trajectory)} 个点...")
        print("绿色线条: 正常轨迹 | 红色线条: 碰撞段")

        trajectory_line_ids = []
        trajectory_points = []

        for step, (joint_angles_deg, x_pos, y_pos) in enumerate(trajectory):
            # 将角度从度转换为弧度
            joint_angles_rad = [round(math.radians(angle),4) for angle in joint_angles_deg]
            # 设置桁轨位置
            p.setJointMotorControl2(
                self.obstacle, self.gantry_prismatic_joints[0],
                p.POSITION_CONTROL, targetPosition=y_pos
            )
            p.setJointMotorControl2(
                self.obstacle, self.gantry_prismatic_joints[1],
                p.POSITION_CONTROL, targetPosition=x_pos
            )

            # 设置机械臂关节
            for i, angle in enumerate(joint_angles_rad[:6]):  # 限制为6个关节
                if i < len(self.robot_joints):
                    p.setJointMotorControl2(
                        self.robot_id, self.robot_joints[i],
                        p.POSITION_CONTROL, targetPosition=angle
                    )

            if gui_mode:
                #等待运行到位
                while (round(self.y_cur_pos,4) != round(y_pos,4)) or (round(self.x_cur_pos,4) != round(x_pos,4) or self.arm_cur_angle != joint_angles_rad):
                    self.x_cur_pos, self.y_cur_pos = self.get_prismatic_joints_pose()
                    self.arm_cur_angle = self.get_arm_joints_angle()
                    # print(f"arm_cur_angle:{self.arm_cur_angle}")
                    # print(f"joint_angles_rad:{joint_angles_rad}")
                    for _ in range(10):
                        p.stepSimulation()
                        if real_time:
                            time.sleep(step_time / 10)
                    collision = self.check_collision()
                    if collision:
                        break
            else:
                if step < 2:
                    while (round(self.y_cur_pos, 4) != round(y_pos, 4)) or (
                            round(self.x_cur_pos, 4) != round(x_pos, 4) or self.arm_cur_angle != joint_angles_rad):
                        self.x_cur_pos, self.y_cur_pos = self.get_prismatic_joints_pose()
                        self.arm_cur_angle = self.get_arm_joints_angle()
                        # print(f"arm_cur_angle:{self.arm_cur_angle}")
                        # print(f"joint_angles_rad:{joint_angles_rad}")
                        for _ in range(10):
                            p.stepSimulation()
                            if real_time:
                                time.sleep(step_time / 10)
                        collision = self.check_collision()
                        if collision:
                            break
                else:
                    for _ in range(10):
                        p.stepSimulation()
                        if real_time:
                            time.sleep(step_time / 10)



            # 获取末端执行器位置
            ee_position, ee_orientation = self.get_end_effector_pose()

            # 检查碰撞
            collision = self.check_collision()
            if collision:
                results['collision_detected'] = True
                results['collision_step'] = step
                results['collision_position'] = ee_position
                results['collision_info'] = self.collision_info
                print(f"步骤 {step}: 检测到碰撞！")

            if gui_mode:
                # 添加轨迹点可视化
                # if collision:
                #     # 红色球体表示碰撞点
                #     sphere_visual = p.createVisualShape(
                #         p.GEOM_SPHERE, radius=0.03,
                #         rgbaColor=[1, 0, 0, 0.8]
                #     )
                # else:
                #     # 蓝色小球表示正常轨迹点
                #     sphere_visual = p.createVisualShape(
                #         p.GEOM_SPHERE, radius=0.015,
                #         rgbaColor=[0, 0.5, 1, 0.6]
                #     )
                #
                # sphere_body = p.createMultiBody(
                #     baseMass=0, baseVisualShapeIndex=sphere_visual,
                #     basePosition=ee_position
                # )
                # trajectory_line_ids.append(sphere_body)

                # 绘制轨迹线段
                if len(trajectory_points) > 0:
                    start_pos = trajectory_points[-1]
                    end_pos = ee_position

                    if collision:
                        # 红色线段表示碰撞段
                        line_id = p.addUserDebugLine(
                            start_pos, end_pos,
                            lineColorRGB=[1, 0, 0],
                            lineWidth=5
                        )
                    else:
                        # 绿色线段表示正常轨迹
                        line_id = p.addUserDebugLine(
                            start_pos, end_pos,
                            lineColorRGB=[0, 1, 0.2],
                            lineWidth=3
                        )

                    trajectory_line_ids.append(line_id)

                trajectory_points.append(ee_position)

            if collision:
                break
        print(f"轨迹播放完成。{'检测到碰撞' if results['collision_detected'] else '无碰撞'}")
        return results



    def cleanup(self):
        """清理仿真环境"""
        p.disconnect(self.physics_client)
        print("仿真环境已清理")





def load_trajectory_from_file(trajectory_file=" "):
    """从文件读取轨迹数据"""
    print(f"正在读取轨迹文件: {trajectory_file}")

    try:
        trajectory = []
        with open(trajectory_file, 'r', encoding='utf-8') as file:
            lines = file.readlines()

        print(f"文件包含 {len(lines)} 行数据")

        # 分析第一行数据来确定格式
        if lines:
            first_line = lines[0].strip()
            print(f"第一行原始数据: {first_line}")

            # 更仔细地解析数据，确保正确处理负号
            raw_values = first_line.split(',')
            print(f"分割后的原始值: {raw_values[:10]}")  # 显示前10个原始值

            try:
                sample_data = []
                for val in raw_values:
                    val = val.strip()  # 去除空格
                    if val:  # 确保不是空字符串
                        sample_data.append(float(val))

                print(f"样本数据长度: {len(sample_data)}")
                print(f"样本数据前10个值: {sample_data[:10]}")

                # 检查是否有负数
                negative_count = sum(1 for x in sample_data[:10] if x < 0)
                print(f"前10个值中的负数个数: {negative_count}")

            except ValueError as e:
                print(f"解析第一行数据时出错: {e}")

        for line_num, line in enumerate(lines):
            line = line.strip()
            if not line:  # 跳过空行
                continue

            try:
                # 更安全的数据解析方式
                raw_values = line.split(',')
                data = []

                for val in raw_values:
                    val = val.strip()  # 去除前后空格
                    if val:  # 确保不是空字符串
                        try:
                            # 直接转换为float，保持负号
                            float_val = float(val)
                            data.append(float_val)
                        except ValueError:
                            print(f"警告: 第 {line_num + 1} 行包含无效数值: '{val}'")
                            continue

                if len(data) < 6:  # 至少需要6个关节角度
                    print(f"警告: 第 {line_num + 1} 行有效数据不足({len(data)}个)，跳过")
                    continue

                # 根据数据格式提取关节角度和桁轨位置
                # 前6个是机械臂关节角度（度）
                joint_angles_deg = data[0:6]

                # 桁轨位置 - 根据数据长度灵活处理
                if len(data) >= 8:
                    # 如果有足够数据，使用第7、8个值作为桁轨位置
                    gantry_x = data[7] / 1000.0  # 转换为米
                    gantry_y = data[6] / 1000.0  # 转换为米
                elif len(data) >= 7:
                    # 如果只有7个数据，第7个作为Y轴，X轴设为0
                    gantry_x = 0.0
                    gantry_y = data[6] / 1000.0
                else:
                    # 如果只有6个关节角度，桁轨位置设为0
                    gantry_x = 0.0
                    gantry_y = 0.0

                # 将角度从度转换为弧度（保持负数）
                # joint_angles_rad = [math.radians(angle) for angle in joint_angles_deg]

                # 格式：(joint_angles, x_pos, y_pos)
                trajectory.append((joint_angles_deg, gantry_x, gantry_y))  #此处关节角是角度制


            except (ValueError, IndexError) as e:
                print(f"警告: 第 {line_num + 1} 行数据格式错误，跳过: {e}")
                print(f"  原始数据: {line[:100]}")  # 显示前100个字符
                continue

        print(f"成功读取 {len(trajectory)} 个轨迹点")

        if len(trajectory) == 0:
            print("错误: 没有有效的轨迹数据")
            return None

        # 显示前几个轨迹点的信息，重点检查负数
        print("\n前3个轨迹点预览（检查负数）:")
        for i in range(min(3, len(trajectory))):
            joints, x_pos, y_pos = trajectory[i]
            print(f"点 {i + 1}:")
            print(f"  关节角度(度): {[round(j, 2) for j in joints]}")
            print(f"  桁轨X位置(米): {round(x_pos, 3)}")
            print(f"  桁轨Y位置(米): {round(y_pos, 3)}")

        return trajectory

    except FileNotFoundError:
        print(f"错误: 找不到文件 {trajectory_file}")
        print("请确保文件在当前目录下")
        return None

    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        import traceback
        traceback.print_exc()
        return None


def run_trajectory_simulation(trajectory, step_time=0,gui_mode: bool = True):
    """
        运行轨迹仿真
        Args:
            trajectory: 轨迹列表，每个元素为(joint_angles, x_pos, y_pos)
            gui_mode : 是否可视化
            step_time: 每步的时间间隔
            real_time: 是否实时显示

    """
    if trajectory is None:
        print("错误: 没有有效的轨迹数据，无法运行仿真")
        return None

    print("启动轨迹测试仿真...")
    simulator = InteractiveRobotTrajectorySimulator(gui_mode)

    try:

        print(f"开始播放轨迹，步进时间: {step_time}s")

        # 执行轨迹
        results = simulator.simulate_trajectory(trajectory, gui_mode,step_time=step_time, real_time=True)

        # 分析结果
        print(f"\n=== 轨迹测试结果 ===")
        print(f"碰撞检测: {'发生碰撞' if results['collision_detected'] else '无碰撞'}")

        if results['collision_detected']:
            print(f"碰撞步骤: {results['collision_step']}")

        return results

    except Exception as e:
        print(f"运行轨迹仿真时发生错误: {e}")
        import traceback
        traceback.print_exc()
        return None


def trajectory_test_example():
    """轨迹测试主函数"""
    # 第一步：读取轨迹文件
    trajectory = load_trajectory_from_file("inverse_solutions_new3(2).txt")

    res = run_trajectory_simulation(trajectory, step_time=0, gui_mode = 1)

    return res

if __name__ == "__main__":

    res = trajectory_test_example()
    if(res['collision_detected']):
        print(f"{res}")
    while True:
        p.stepSimulation()
